#!/usr/bin/python
'''
Marks parts of contig within 3k of edge that have greater than x coverage.
'''
import sys
import numpy as np
import h5py

# parameters.
pileup_file = sys.argv[1]
out_file = sys.argv[2]
cov = float(sys.argv[3])
edge = 3000
window = 50

# declarations.
names_dt = np.dtype([\
		('name', np.str_, 200),\
		('start', np.int),\
		('stop', np.int),\
		])
		
profile_dt = np.int


########
# open hdf5 files.
print "indexing h5 data."
h5_in = h5py.File(pileup_file, 'r')
#names = h5_in['names']
#profiles = h5_in['profiles']
names = h5_in['names'][0::]
profiles = h5_in['profiles'][0::]

# build id to index dictionary.
ctg_to_idx = {}
for i in range(len(names)): 
	ctg_to_idx[names[i][0]] = i
	#if i > 10000: break
 

# Loop over edges.
print "Searching for high coverage windows."
fout = open(out_file, "w")
for ctg in ctg_to_idx:
	# Get info.
	if ctg not in ctg_to_idx: continue
	
	idx = ctg_to_idx[ctg]
	start = names[idx][1]
	stop = names[idx][2]

	# Just check whole contig once.
	if stop + start < 2 * edge:
		continue
	
	
	# loop over 5' edge.
	run_start = -2
	run_stop = -5
	for i in range(start,start + edge - window):
		# calc window.
		tmp = profiles[i:i+window]
		avg = np.average(tmp)
		
		# see if we start or continue a run.
		if avg >= cov:
			# see if we start.
			if run_start == run_stop:
				# set start.
				#print "start", i, avg
				run_start = i
			
			# extend stop.
			run_stop = i + window
			#print "extend", i, profiles[i], avg, run_start, run_stop
		else:
			# process finished run.
			if run_stop >= 0 and run_start >= 0 and run_stop - run_start > 1:
				# print filter.
				wavg = np.average(profiles[run_start:run_stop])
				cstart = run_start - start
				cstop = run_stop - start
				fout.write("%f\t%s\t%i\t%i\n" % (wavg, ctg, cstart, cstop))

			#print "nothing", i, avg
			run_start = run_stop = i


		
# Close h5 connection.
fout.close()
h5_in.close()
	





